{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyOxCMAyx5VokSWvYs8NMPH3",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/LC1332/Chat-Haruhi-Suzumiya/blob/main/notebook/ChatHaruhi%E7%9A%84%E5%B0%8F%E8%AF%B4%E6%95%B0%E6%8D%AE%E7%94%9F%E6%88%90.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "- [x] 统计换行数量\n",
        "- [x] 统计引号数量\n",
        "- [ ] prompt改写成多样化的prompt"
      ],
      "metadata": {
        "id": "S024bFjfrRs1"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0Z_hevNCrRIJ",
        "outputId": "53ed9cb8-19ff-4724-e96f-4bed2efa324c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mounted at /content/drive\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8_c1yqgApHy3",
        "outputId": "bdb811aa-5e20-4e63-9e9c-56cb16f41bcc"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "part_0.txt   part_1.txt  part_3.txt  part_5.txt  part_7.txt  part_9.txt\n",
            "part_10.txt  part_2.txt  part_4.txt  part_6.txt  part_8.txt\n"
          ]
        }
      ],
      "source": [
        "!ls /content/drive/MyDrive/Wuxia/input/all_direct"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "fname = \"/content/drive/MyDrive/Wuxia/input/all_direct/part_10.txt\"\n",
        "\n",
        "import json\n",
        "\n",
        "datas = []\n",
        "\n",
        "with open(fname, \"r\", encoding=\"utf-8\") as f:\n",
        "    for line in f:\n",
        "        if line.strip() == \"\":\n",
        "            continue\n",
        "        data = json.loads(line)\n",
        "        datas.append(data)"
      ],
      "metadata": {
        "id": "aOGYrFf0sLrA"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(len(datas))\n",
        "print(datas[0].keys())"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "-e3-cQucsXgN",
        "outputId": "d7501d53-c322-4d9f-fbb6-cbc74f2a7c54"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "20205\n",
            "dict_keys(['source', 'text'])\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def count_quotes(text):\n",
        "\n",
        "    quote_chars = '''「」\"“”'''\n",
        "    count = 0\n",
        "    for char in text:\n",
        "        if char in quote_chars:\n",
        "            count += 1\n",
        "    return count"
      ],
      "metadata": {
        "id": "Tjq9GTNssyfc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def count_paragraphs(text):\n",
        "    lines = text.split(\"\\n\")\n",
        "    count = 0\n",
        "    for line in lines:\n",
        "        if line.strip() != \"\":\n",
        "            count += 1\n",
        "    return count"
      ],
      "metadata": {
        "id": "-yIyDBqvtQTy"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from tqdm import tqdm\n",
        "n_paragraphs = []\n",
        "n_quotes = []\n",
        "\n",
        "for data in tqdm(datas):\n",
        "    n_paragraphs.append(count_paragraphs(data[\"text\"]))\n",
        "    n_quotes.append(count_quotes(data[\"text\"]))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "BdNMMR--saEL",
        "outputId": "714dec61-b402-4932-f8e5-6ab7c93cbaf2"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 20205/20205 [00:06<00:00, 3006.08it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "对于一个字符串line，检查里面的'“'或者'\"'或者'”'出现位置。并把所有位置合并为一个pos的list，如果pos的长度为偶数，则把奇数位置替换为「把偶数位置替换为」，请用python为我实现"
      ],
      "metadata": {
        "id": "sAFytaQowe1n"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def change_quotes(text):\n",
        "    lines = text.split(\"\\n\")\n",
        "\n",
        "    quotes = ['“', '”', '\"']\n",
        "\n",
        "    new_text = \"\"\n",
        "    for line in lines:\n",
        "        if line.strip() == \"\":\n",
        "            continue\n",
        "        positions = []\n",
        "        for i, char in enumerate(line):\n",
        "            if char in quotes:\n",
        "                positions.append(i)\n",
        "        if len(positions) % 2 == 0:\n",
        "            for i in range(len(positions)):\n",
        "                if i % 2 != 0:\n",
        "                    line = line[:positions[i]] + '」' + line[positions[i]+1:]\n",
        "                else:\n",
        "                    line = line[:positions[i]] + '「' + line[positions[i]+1:]\n",
        "        new_text += line + \"\\n\"\n",
        "    return new_text\n"
      ],
      "metadata": {
        "id": "BZIjnTDVwZL4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "prompt_templates = \"\"\"Proceed with the paragraph, maintaining the same style of language.\n",
        "Keep writing the paragraph, using a similar linguistic style.\n",
        "Continue the paragraph, ensuring the language style remains consistent.\n",
        "Go on with the paragraph in a similar style of language.\n",
        "Carry on writing the paragraph, sticking to the current language style.\n",
        "Extend the paragraph, keeping the language tone and style similar.\n",
        "Resume the paragraph, mirroring the existing language style.\n",
        "Continue crafting the paragraph, maintaining the established language style.\n",
        "Keep extending the paragraph, following the same style of language.\n",
        "Proceed with extending the paragraph, preserving the linguistic style.\n",
        "Please compose a paragraph maintaining a uniform style throughout.\n",
        "Write a paragraph ensuring that the style is consistent from start to finish.\n",
        "Your assignment is to craft a paragraph with a consistent stylistic approach.\n",
        "Create a paragraph where the style remains constant throughout.\n",
        "Construct a paragraph, keeping the writing style uniform in the entire passage.\n",
        "Develop a paragraph that demonstrates consistency in style.\n",
        "Focus on writing a paragraph with a steady and unvarying style.\n",
        "Your challenge is to author a paragraph that maintains a consistent style.\n",
        "Compose a paragraph, ensuring stylistic consistency throughout.\n",
        "Produce a paragraph in which the style does not vary from beginning to end.\n",
        "Your task is to write paragraph in a consistent style\n",
        "对于下面的Paragraph，使用一致性的风格进行续写\n",
        "请继续下面的段落，并保持原有的写作风格一致。\n",
        "在接下来的写作中，延续这一段落的风格，并进行扩展。\n",
        "针对以下段落，以相同的风格进行补充写作。\n",
        "续写下面的段落，并确保风格与原文保持一致。\n",
        "请以与下面段落相同的风格，进行创作性的续写。\n",
        "保持原段落的写作风格，对其进行扩展和深化。\n",
        "延续以下段落的风格，进行相应的续写工作。\n",
        "请在下面的段落后续写，确保风格的连贯性。\n",
        "继续这一段落的故事，注意保持原有风格的一致性。\n",
        "针对下文，以一种风格上连贯的方式进行创作性扩写。\"\"\"\n",
        "\n",
        "prompt_templates = prompt_templates.split(\"\\n\")\n",
        "\n",
        "for i in range(len(prompt_templates)):\n",
        "    prompt_templates[i] = prompt_templates[i].strip()"
      ],
      "metadata": {
        "id": "KRIPcXfDypSY"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q tiktoken\n",
        "import tiktoken\n",
        "\n",
        "enc = tiktoken.get_encoding(\"cl100k_base\")"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CoS9e1Xi7wxh",
        "outputId": "4415706e-cf29-471a-a320-becfd7fdaadd"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "\u001b[2K     \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m2.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
            "\u001b[?25h\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "llmx 0.0.15a0 requires cohere, which is not installed.\n",
            "llmx 0.0.15a0 requires openai, which is not installed.\u001b[0m\u001b[31m\n",
            "\u001b[0m"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title 定义strong_divide\n",
        "def divide_str(s, sep=['\\n', '.', '。'], consider_quote = True ):\n",
        "    mid_len = len(s) // 2  # 中心点位置\n",
        "    best_sep_pos = len(s) + 1  # 最接近中心点的分隔符位置\n",
        "    best_sep = None  # 最接近中心点的分隔符\n",
        "    for curr_sep in sep:\n",
        "        sep_pos = s.rfind(curr_sep, 0, mid_len + mid_len - 2)  # 从中心点往左找分隔符\n",
        "        if sep_pos > 0 and abs(sep_pos - mid_len) < abs(best_sep_pos -\n",
        "                                                        mid_len):\n",
        "            if consider_quote == True:\n",
        "                if s[sep_pos + 1] in ['\"', \"'\",\"“\",\"」\"]:\n",
        "                    sep_pos += 1\n",
        "\n",
        "            best_sep_pos = sep_pos\n",
        "            best_sep = curr_sep\n",
        "    if not best_sep:  # 没有找到分隔符\n",
        "        return s, ''\n",
        "    return s[:best_sep_pos + 1], s[best_sep_pos + 1:]\n",
        "\n",
        "\n",
        "def strong_divide(s):\n",
        "    # left, right = divide_str(s)\n",
        "\n",
        "    # if right != '':\n",
        "    #     return left, right\n",
        "\n",
        "    whole_sep = ['\\n', '.', '，', '、', ';', ',', '；',\\\n",
        "                 '：', '！', '？', '(', ')', '”', '“', \\\n",
        "                 '’', '‘', '[', ']', '{', '}', '<', '>', \\\n",
        "                 '/', '''\\''', '|', '-', '=', '+', '*', '%', \\\n",
        "               '$', '''#''', '@', '&', '^', '_', '`', '~',\\\n",
        "                 '·', '…']\n",
        "    left, right = divide_str(s, sep=whole_sep)\n",
        "\n",
        "    if right != '':\n",
        "        return left, right\n",
        "\n",
        "    mid_len = len(s) // 2\n",
        "    return s[:mid_len], s[mid_len:]\n",
        "\n",
        "def super_divide( s, max_len_token = 300, iter_deep = 1):\n",
        "    n_token = len(enc.encode(s))\n",
        "    if n_token <= max_len_token or iter_deep > 2:\n",
        "        return [s]\n",
        "\n",
        "    left, right = strong_divide(s)\n",
        "    return super_divide(left, max_len_token, iter_deep+1) + \\\n",
        "            super_divide(right, max_len_token, iter_deep+1)"
      ],
      "metadata": {
        "id": "FbaYTY5e7pMT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def segment_paragraph(text):\n",
        "    raw_lines = text.split(\"\\n\")\n",
        "    lines = []\n",
        "\n",
        "    all_token_count = 0\n",
        "\n",
        "    for line in raw_lines:\n",
        "        line = line.strip()\n",
        "        if line == \"\":\n",
        "            continue\n",
        "        n_token = len(enc.encode(line))\n",
        "        all_token_count += n_token\n",
        "\n",
        "        if n_token > 600:\n",
        "            segments = super_divide(line)\n",
        "            lines.extend(segments)\n",
        "        else:\n",
        "            lines.append(line)\n",
        "\n",
        "    seg_token = all_token_count // 3\n",
        "\n",
        "    if seg_token > 1500:\n",
        "        seg_token = 1500\n",
        "\n",
        "    current_seg_token = 0\n",
        "    current_seg = \"\"\n",
        "\n",
        "    segments = []\n",
        "    for line in lines:\n",
        "        line = line.strip('\"#-。，, \\n')\n",
        "        if line.strip() == \"\":\n",
        "            continue\n",
        "        n_token = len(enc.encode(line))\n",
        "        if current_seg_token + n_token > seg_token:\n",
        "            if current_seg_token > 0:\n",
        "                segments.append(current_seg.strip())\n",
        "                current_seg = line\n",
        "                current_seg_token = n_token\n",
        "        else:\n",
        "            current_seg += \"\\n\" + line\n",
        "            current_seg_token += n_token\n",
        "\n",
        "    if current_seg_token > 0:\n",
        "        segments.append(current_seg.strip())\n",
        "\n",
        "    return segments"
      ],
      "metadata": {
        "id": "qJ8eQujTz44g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(10//3)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OUnG6K_l0JVc",
        "outputId": "e09e34b3-21e9-4746-fb48-0c1ba7cc72c9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "3\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "source_pool = set()"
      ],
      "metadata": {
        "id": "sKvwEjNt1e61"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def find_first_sep( text ):\n",
        "    first_sep = 9999\n",
        "    first_sep_second_choice = 9999\n",
        "\n",
        "    find_first_ket = text[:50].find(\"」\",0)\n",
        "\n",
        "    if find_first_ket > 0:\n",
        "        return find_first_ket\n",
        "\n",
        "    if len(text) <= 50:\n",
        "        return len(text)\n",
        "    for i, c in enumerate(text):\n",
        "        if i > 50:\n",
        "            break\n",
        "        if c in ['。','\\n']:\n",
        "            first_sep = min(first_sep, i)\n",
        "        elif c in ['\\n', '.', '，', '、', ';', ',', '；',\\\n",
        "                 '：', '！', '？', '(', ')', '”', '“', \\\n",
        "                 '’', '‘', '[', ']', '{', '}', '<', '>', \\\n",
        "                 '/', '''\\''', '|', '-', '=', '+', '*', '%', \\\n",
        "               '$', '''#''', '@', '&', '^', '_', '`', '~',\\\n",
        "                 '·', '…']:\n",
        "            first_sep_second_choice = min(first_sep_second_choice, i)\n",
        "\n",
        "    for i in range(5):\n",
        "        if len(text) > first_sep and text[first_sep+1] in ['\\n', '.', '，', '、', ';', ',', '；',\\\n",
        "                    '：', '！', '？', '(', ')', '”', '“', \\\n",
        "                    '’', '‘', '[', ']', '{', '}', '<', '>', \\\n",
        "                    '/', '''\\''', '|', '-', '=', '+', '*', '%', \\\n",
        "                '$', '''#''', '@', '&', '^', '_', '`', '~',\\\n",
        "                    '·', '…','「','」']:\n",
        "            first_sep += 1\n",
        "\n",
        "    if first_sep <= 50:\n",
        "        return first_sep\n",
        "    elif first_sep_second_choice <= 50:\n",
        "        return first_sep_second_choice\n",
        "    else:\n",
        "        return 50"
      ],
      "metadata": {
        "id": "huG54_dK-NXE"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "count = 0\n",
        "\n",
        "import random\n",
        "\n",
        "save_datas = []\n",
        "\n",
        "for i in tqdm(range(len(datas))):\n",
        "    if n_quotes[i] < 8 or n_paragraphs[i] < 9:\n",
        "        continue\n",
        "\n",
        "    source = datas[i][\"source\"]\n",
        "    source_pool.add(source)\n",
        "\n",
        "    text = datas[i][\"text\"]\n",
        "    text = change_quotes( text )\n",
        "    segments = segment_paragraph( text )\n",
        "\n",
        "    context = random.choice(prompt_templates)\n",
        "\n",
        "    context += \"\\nParagraph:\\n###\\n\" + segments[0] + \"\\n###\\n\"\n",
        "\n",
        "    first_sep = find_first_sep(segments[1])\n",
        "\n",
        "    context += segments[1][:first_sep+1]\n",
        "\n",
        "    target = segments[1][first_sep+1:]\n",
        "\n",
        "    for j in range(2, len(segments)):\n",
        "        target += \"\\n\\n\" + segments[j]\n",
        "\n",
        "    save_data = {\n",
        "        \"context\": context,\n",
        "        \"target\": target\n",
        "    }\n",
        "\n",
        "    save_datas.append(save_data)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "gxHPWRf_wQOd",
        "outputId": "0286d98e-c033-4a8d-ba89-25061295f94c"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 20205/20205 [01:41<00:00, 199.99it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "save_name = \"/content/classic_novel_writing_18k.jsonl\"\n",
        "\n",
        "with open(save_name, 'w', encoding='utf8') as f:\n",
        "    for data in save_datas:\n",
        "        json.dump(data, f, ensure_ascii=False)\n",
        "        f.write('\\n')"
      ],
      "metadata": {
        "id": "UJDMmlQg0ulC"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(context)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "VNT79sWMtnQ1",
        "outputId": "9fb65af3-dc95-4a4f-b54c-759f5c566bc9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Your challenge is to author a paragraph that maintains a consistent style.\n",
            "Paragraph:\n",
            "###\n",
            "八月末的一天清早，白嘉轩起来洗脸漱口时，他的冒死破禁而且显出怀孕徵兆的妻子仙草正坐在纺线车前嗡嗡嗡嗡地转动着车把儿，锭子上已经结下一枚茭白大小的白色线穗了。母亲也早已起来，在自个独居的里屋炕上摇转着纺车。他坐在父亲在世时常坐的那把靠背椅子上，喝看酽茶，用父亲死後留下的那把白铜水烟袋过著早瘾。父亲死後，他每天晚上在母亲落枕前和清早起床後都到里屋里坐一会儿。两架纺车嗡嗡吱吱的声音互相衔接，互相重合，此声间歇，彼声响起，把沉稳和谐的气氛弥漫到四合院的每一个角落。白嘉轩沉浸在这古老悠远而又新鲜活泼的乐曲里，浑身的筋骨和血液就鼓涨起来\n",
            "长工鹿三把犁铧套绳收拾齐备，从马号里牵出红马拴在院子里的石雕拴马桩上，扯着大步走进院庭，大声询问种子的事。嘉轩从里屋走出来：「你先喝口茶。」鹿三站在院庭里说他不喝，仍然询问麦子和豌豆掺和的比例，二八还是三七？嘉轩说：「这块地种药材。种子你甭管，我拿着。」说着喷出一口烟，吹净水烟筒里的烟灰，放下水烟壶，喝下最後一盅茶，就赳赳地走出街门，进入马号。鹿三解下红马牵着，套上犁杖。嘉轩扛起沉重的铁齿大耙子，腋下挟着一把镢头和一把竹条扫帚。鹿三回过头问：「你拿扫帚做啥？」嘉轩也不解释：「拿就是有用嘛。」鹿三就不再问。主仆二人走过街巷，出了村子，走下河滩，红马拖着空犁在田间土路上撞出瞠瞠瞠的声响\n",
            "田野已经改换过另一种姿容，斑斓驳杂的秋天的色彩像羽毛一样脱光褪尽荡然无存了，河川里呈现出一种喧闹之後的沉静。灌渠渠沿和井台上堆积着刚刚从田地里清除出来的包谷秆子。麦子播种几近尾声，刚刚播种不久的田块裸露着湿漉漉的泥土，早种的田地已经泛出麦苗幼叶的嫩绿。秋天的淫雨季节已告结束，长久弥漫在河川和村庄上空的阴霾和沉闷已全部廓清。大地简洁而素雅，天空开阔而深远。清晨的冷气使人精神抖擞\n",
            "###\n",
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "# 额外20k数据的获取\n",
        "\n",
        "下面从全小说库进行探查，但是这个时候我们只考虑对话较多的数据"
      ],
      "metadata": {
        "id": "_J-eOjxAByim"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "novel2count = {}\n",
        "\n",
        "import json\n",
        "\n",
        "datas = []\n",
        "for part_id in tqdm(range(0, 10)):\n",
        "    fname = \"/content/drive/MyDrive/Wuxia/input/all_direct/part_\"+ str(part_id) +\".txt\"\n",
        "\n",
        "\n",
        "    with open(fname, \"r\", encoding=\"utf-8\") as f:\n",
        "        for line in f:\n",
        "            if line.strip() == \"\":\n",
        "                continue\n",
        "            data = json.loads(line)\n",
        "            source = data[\"source\"]\n",
        "            if source in source_pool:\n",
        "                continue\n",
        "\n",
        "            if source not in novel2count:\n",
        "                novel2count[source] = 0\n",
        "\n",
        "            if novel2count[source] > 20:\n",
        "                continue\n",
        "\n",
        "            n_quote = count_quotes(data[\"text\"])\n",
        "            if n_quote >= 110:\n",
        "                continue\n",
        "\n",
        "            novel2count[source] += 1\n",
        "            datas.append(data)\n",
        "\n",
        "    print(len(datas))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aoNeoQPu83aN",
        "outputId": "fd2a183a-6d89-491d-de13-d979822bfc5f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            " 10%|█         | 1/10 [00:06<00:58,  6.46s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "4949\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 20%|██        | 2/10 [00:17<01:11,  8.97s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "10146\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 30%|███       | 3/10 [00:23<00:55,  7.95s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "15219\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 40%|████      | 4/10 [00:33<00:50,  8.45s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "20310\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 50%|█████     | 5/10 [00:44<00:46,  9.37s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "25175\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 60%|██████    | 6/10 [00:57<00:42, 10.66s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "30168\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 70%|███████   | 7/10 [01:11<00:35, 11.87s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "35231\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 80%|████████  | 8/10 [01:23<00:23, 11.95s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "40409\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\r 90%|█████████ | 9/10 [01:43<00:14, 14.52s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "45464\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 10/10 [01:58<00:00, 11.81s/it]"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "50427\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "count = 0\n",
        "\n",
        "import random\n",
        "\n",
        "save_datas = []\n",
        "\n",
        "for i in tqdm(range(len(datas))):\n",
        "    source = datas[i][\"source\"]\n",
        "    source_pool.add(source)\n",
        "\n",
        "    text = datas[i][\"text\"]\n",
        "    text = change_quotes( text )\n",
        "    segments = segment_paragraph( text )\n",
        "\n",
        "    if len(segments) < 3:\n",
        "        continue\n",
        "\n",
        "    context = random.choice(prompt_templates)\n",
        "\n",
        "    context += \"\\nParagraph:\\n###\\n\" + segments[0] + \"\\n###\\n\"\n",
        "\n",
        "    first_sep = find_first_sep(segments[1])\n",
        "\n",
        "    context += segments[1][:first_sep+1]\n",
        "\n",
        "    target = segments[1][first_sep+1:]\n",
        "\n",
        "    for j in range(2, len(segments)):\n",
        "        target += \"\\n\\n\" + segments[j]\n",
        "\n",
        "    save_data = {\n",
        "        \"context\": context,\n",
        "        \"target\": target\n",
        "    }\n",
        "\n",
        "    save_datas.append(save_data)\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9_VEPw4BF3Uk",
        "outputId": "7b8d8f48-dfd2-4274-8298-e9119745876f"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 50427/50427 [04:56<00:00, 169.95it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "import json\n",
        "\n",
        "save_name = \"/content/drive/MyDrive/Wuxia/input/all_direct/dialogueful_novel.jsonl\"\n",
        "\n",
        "with open(save_name, 'w', encoding='utf8') as f:\n",
        "    for data in save_datas:\n",
        "        json.dump(data, f, ensure_ascii=False)\n",
        "        f.write('\\n')"
      ],
      "metadata": {
        "id": "3IT9wDQdG-Qw"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(len(datas))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JA9DTrmiDa_i",
        "outputId": "180e9137-6088-4e2d-f4c5-26ee945b09c9"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "87620\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for data in tqdm(datas):\n",
        "    n_quotes.append(count_quotes(data[\"text\"]))"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "11Dv4I4GDcn2",
        "outputId": "0c4a1dca-a639-4b29-af90-c63a8ef00f91"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "100%|██████████| 98666/98666 [00:29<00:00, 3312.06it/s]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "n_quotes_sorted = sorted(n_quotes, reverse=True)"
      ],
      "metadata": {
        "id": "3TYsFZ5iDoUG"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(n_quotes_sorted[2000])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "3780EZkXDkcC",
        "outputId": "41484cdd-0668-4715-9bde-d6e5b4368351"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "116\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "print(n_quotes_sorted[2000])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zoK8svzXD0BG",
        "outputId": "4c083016-e00d-42e0-b474-24c66ea19b32"
      },
      "execution_count": null,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "132\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "z1oHyqNsELxR"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}